/*BEGIN_LEGAL 
Intel Open Source License 

Copyright (c) 2002-2005 Intel Corporation 
All rights reserved. 
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.  Redistributions
in binary form must reproduce the above copyright notice, this list of
conditions and the following disclaimer in the documentation and/or
other materials provided with the distribution.  Neither the name of
the Intel Corporation nor the names of its contributors may be used to
endorse or promote products derived from this software without
specific prior written permission.
 
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE INTEL OR
ITS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
END_LEGAL */

//mcfTrace.cpp
//Multithreaded branch, call, and ret instruction tracing

// Bugs and criticisms can be directed to
// myers.albert@gmail.com

//changes highlighted by ++ sign - Amrish K. Tewar

//Basic control tracing for multithreaded programs


#include <unistd.h>
#include <stdio.h>
#include <string>
#include <iostream>
#include <sstream>
#include <iomanip>
#include <fstream>
#include <deque>
#include <limits>
#include <sys/time.h>
#include <sys/stat.h>
#include "pin.H"
#include "instlib.H"
#include "mcfTraceAssist_akt.h"
#include "mcfTrace_akt.h"
//change start
// 1 -#include <vector>


//change stop
//Can be used to filter which parts of the binary to instrument
INSTLIB::FILTER filter;



//instrument instructions
//three seperate instrumentation routines - ASCII, ASCII with disassembly, and binary
VOID ASCIIControlTrace(TRACE trace, VOID *v)
{
   //return if not a selected routin
   if( !filter.SelectTrace(trace) )
      return;
    
   //iterate through instructions
   for(BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl) )
   {
   //UINT32 c = BBL_NumIns(bbl);
      for(INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins) )
      {
         INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)SetFastForwardAndLength, IARG_THREAD_ID, IARG_END);
         //From the pin manual
	 //HasFallThrough is TRUE for instructions which don't change the control flow (most instructions), 
	 //or for conditional branches (which might change the control flow, but might not),
	 //
	 //HasFallThrough is FALSE for unconditional branches and calls 
	 //(where the next instruction to be executed is always explicitly specified).
             
         // Is Unconditional and Direct and Taken
	 if( INS_IsDirectBranchOrCall(ins) && !INS_HasFallThrough(ins) )
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)Emit_UnconditionalDirect_ASCII, 
                        //Args: Thread ID, Instruction Address, Target Address, # of instruction in BBL // 2 -  IARG_UINT32, c,
                           IARG_THREAD_ID, IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, 
						   IARG_END);

         // Is Conditional and Direct
         else if( INS_IsDirectBranchOrCall(ins) && INS_HasFallThrough(ins) )             
                 INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)Emit_ConditionalDirect_ASCII, 
                          //Args: Thread ID, Instruction Address, Target Address, Taken? // 2- IARG_UINT32, c,
                             IARG_THREAD_ID, IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, IARG_BRANCH_TAKEN,  IARG_END);

         // Is Unconditional and Indirect - Returns are indirect
         else if( INS_IsIndirectBranchOrCall(ins) || INS_IsRet(ins) )
                 INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)Emit_UnconditionalIndirect_ASCII, 
                          //Args: Thread ID, Instruction Address, Target Address  //2 -  IARG_UINT32, c,
                             IARG_THREAD_ID, IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR,   IARG_END);
      }
   }
}


VOID DisControlTrace(TRACE trace, VOID *v)
{
   //return if not a selected routine
   if( !filter.SelectTrace(trace) )
      return;
    
   //iterate through instructions
   for(BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl) )
   {
      for(INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins) )
      {
         INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)SetFastForwardAndLength, IARG_THREAD_ID, IARG_END);
         // Is Unconditional and Direct and Taken
         if( INS_IsDirectBranchOrCall(ins) && !INS_HasFallThrough(ins) )
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)Emit_UnconditionalDirect_ASCII_Dis, 
                        //Args: Thread ID, Instruction Address, Target Address, instruction
                           IARG_THREAD_ID, IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, IARG_END);

         // Is Conditional and Direct
         else if( INS_IsDirectBranchOrCall(ins) && INS_HasFallThrough(ins) )
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)Emit_ConditionalDirect_ASCII_Dis, 
                        //Args: Thread ID, Instruction Address, Target Address, Taken?
                           IARG_THREAD_ID, IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, IARG_BRANCH_TAKEN, IARG_END);

         // Is Unconditional and Indirect
         else if( INS_IsIndirectBranchOrCall(ins) || INS_IsRet(ins) )
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)Emit_UnconditionalIndirect_ASCII_Dis, 
                        //Args: Thread ID, Instruction Address, Target Address, instruction
                           IARG_THREAD_ID, IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, IARG_END);
         
      }
   }
}


VOID BinaryControlTrace(TRACE trace, VOID *v)
{
   //return if not a selected routine
   if( !filter.SelectTrace(trace) )
      return;
    
   //iterate through instruction
   for(BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl) )
   {
      for(INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins) )
      {
         INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)SetFastForwardAndLength, IARG_THREAD_ID, IARG_END);
         // Is Unconditional and Direct and Taken
	 if( INS_IsDirectBranchOrCall(ins) && !INS_HasFallThrough(ins) )
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)Emit_UnconditionalDirect_Bin, 
                        //Args: Thread ID, Instruction Address, Target Address 
                           IARG_THREAD_ID, IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, IARG_END);

         // Is Conditional and Direct
         else if( INS_IsDirectBranchOrCall(ins) && INS_HasFallThrough(ins) )
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)Emit_ConditionalDirect_Bin, 
                        //Args: Thread ID, Instruction Address, Target Address, Taken
                           IARG_THREAD_ID, IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, IARG_BRANCH_TAKEN, IARG_END);

         // Is Unconditional and Indirect
         else if( INS_IsIndirectBranchOrCall(ins) || INS_IsRet(ins) )
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)Emit_UnconditionalIndirect_Bin, 
                        //Args: Thread ID, Instruction Address, Target Address 
                           IARG_THREAD_ID, IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, IARG_END);
      }
   }
}


//Called when application closes
VOID Fini(INT32 code, VOID *v)
{
   gettimeofday(&t2, NULL);
   //write remaining descriptors to file
   if(ASCII.Value())
   {
      if(usingCompression)
         WriteASCIIDescriptorTableToCompressor(PIN_GetTid()+1);
      else
         WriteASCIIDescriptorTableToFile(PIN_GetTid());
   }
   else
   {
      if(usingCompression)
         WriteBinaryDescriptorTableToCompressor(PIN_GetTid());
      else
         WriteBinaryDescriptorTableToFile(PIN_GetTid());
   }
   //close 
   closeOutput();
   PrintStatistics();
}


//Called if Pin detaches from application
VOID DetachCallback(VOID *args)
{
   gettimeofday(&t2, NULL);
   //write remaining descriptors to file
   if(ASCII.Value())
   {
      if(usingCompression)
         WriteASCIIDescriptorTableToCompressor(PIN_GetTid()+1);
      else
         WriteASCIIDescriptorTableToFile(PIN_GetTid());
   }
   else
   {
      if(usingCompression)
         WriteBinaryDescriptorTableToCompressor(PIN_GetTid());
      else
         WriteBinaryDescriptorTableToFile(PIN_GetTid());
   }

   //close file/pipe
   closeOutput();
   std::cout << "mcfTrace: Detaching..." << std::endl;

   PrintStatistics();
}

//called on thread creation
/*VOID ThreadStart(THREADID threadid, CONTEXT *ctxt, INT32 flags, VOID *v)
{
   PIN_GetLock(&cout_lock, threadid+1);
   UINT8* tid = new UINT8;
   *tid = numThreads;
   
   //This PIN tool spawns its own thread to write to file, which interrupts the logical ordering of threads
   //Use TLS to record application thread ids
   PIN_SetThreadData(tls_key, tid, threadid);
  //change start 
	 counter_inst.push_back(0); // 1 -
	//thread_data_t* tdata = new thread_data_t;
   //PIN_SetThreadData(tls_key1, tdata, threadid);
  // change stop 
   std::cout << "mcfTrace: thread begin " << static_cast<UINT16>(numThreads) << " " << PIN_GetTid() << std::endl;
   numThreads++;
   PIN_ReleaseLock(&cout_lock);
}*/

VOID ThreadStart(THREADID threadid, CONTEXT *ctxt, INT32 flags, VOID *v)
{
   PIN_GetLock(&cout_lock, threadid+1);
   thread_data_t* tdata = new thread_data_t; //C - 3
   tdata->_tid = numThreads; // C - 3
   tdata->_count = 0; // C - 3
   
   //This PIN tool spawns its own thread to write to file, which interrupts the logical ordering of threads
   //Use TLS to record application thread ids
   PIN_SetThreadData(tls_key, tdata, threadid);
  //change start 
	 //counter_inst.push_back(0); // 1 -
	//thread_data_t* tdata = new thread_data_t;
   //PIN_SetThreadData(tls_key1, tdata, threadid);
  // change stop 
   std::cout << "mcfTrace: thread begin " << static_cast<UINT16>(numThreads) << " " << PIN_GetTid() << std::endl;
   numThreads++;
   PIN_ReleaseLock(&cout_lock);
}


//delete thread local storage
/*VOID ThreadFini(THREADID threadid, const CONTEXT *ctxt, INT32 flags, VOID *v)
{
   delete static_cast<UINT8*>(PIN_GetThreadData(tls_key, threadid));
   //change start
   //delete static_cast<thread_data_t*>(PIN_GetThreadData(tls_key1, threadid));
   //change stop
}*/

//delete thread local storage
VOID ThreadFini(THREADID threadid, const CONTEXT *ctxt, INT32 flags, VOID *v)
{
   cout << "extra...." << endl;
   /*for (INT32 t = 0; t<numThreads; t++)
   {
	thread_data_t* tdata = get_tls(t);
	cout << "what i find is : " << tdata->_count << endl;
   }*/
   delete static_cast<thread_data_t*>(PIN_GetThreadData(tls_key, threadid));
   //change start
   //delete static_cast<thread_data_t*>(PIN_GetThreadData(tls_key1, threadid));
   //change stop
}

INT32 Usage()
{
   //write out usage 
   std::cerr << KnobOutputFile.StringKnobSummary() << endl;
   return -1;
}


int main(int argc, char * argv[])
{
   // Initialize pin
   PIN_InitSymbols();
   //if correct parameters, show usage
   if (PIN_Init(argc, argv)) return Usage();
   
   if(KnobOutputFile.Value() == "mcfTrace_akt.out")
   {
      //Create time stamp for output file so user won't overwrite previous files
      time_t t = time(NULL);
      stamp = *localtime(&t);
      string temp;
      temp = std::to_string(static_cast<long long>(stamp.tm_year + 1900)) + "_" + std::to_string(static_cast<long long>(stamp.tm_mon + 1))+ "_" 
             + std::to_string(static_cast<long long>(stamp.tm_mday)) + "_" + std::to_string(static_cast<long long>(stamp.tm_hour)) + "." 
             + std::to_string(static_cast<long long>(stamp.tm_min)) + "." + std::to_string(static_cast<long long>(stamp.tm_sec));
      //Concatenate file name and time stamp
     
      fileName = KnobOutputFile.Value() + temp;
   }
   else
      fileName = KnobOutputFile.Value();
   //Initialize locks
   PIN_InitLock(&table_lock); 
   PIN_InitLock(&decoder_lock); 
   PIN_InitLock(&cout_lock); 
   PIN_InitLock(&count_lock); 
   //change start
   PIN_InitLock(&count_ins_lock);
   // - 1 PIN_InitLock(&count_ins_lock); 
	//change stop
   //The number of instructions to trace 
   if(Length.Value()==0)
      noLength = TRUE;
   else
      noLength = FALSE; 
    
   //If writing descriptors to binary file
   if( ASCII.Value() == FALSE )
   {
      if( AnnotateDisassembly.Value() == TRUE )
      {
         std::cout << "mcfTrace: Can only annotate disassembly if writing to ASCII file" << std::endl;
         std::exit(-1);
      }

      std::cout << "mcfTrace: Writing to binary file: " << fileName + ".bin" << std::endl  
                << "mcfTrace Descriptor encoding: "
              	<< "ThreadID ( " << sizeof(UINT8) << " byte ), " 
                << "Instruction address ( "  << sizeof(ADDRINT) << " bytes ), " 
	        << "Branch target ( " << sizeof(ADDRINT) << " bytes ), " 
	        << "Branch type and outcome ( " << sizeof(UINT8) << " byte ) " << std::endl; 
   }
   //If writing descriptors to ASCII file	       
   else
   {
      std::cout << "mcfTrace: Writing to text file: " << fileName + ".txt" << std::endl;
      std::cout << "mcfTrace descriptor: tid, instruction address, target address, conditional or unconditional, direct or indirect, outcome" << std::endl;
      //initialize decoder used to get instruction string
      if(AnnotateDisassembly.Value() == TRUE)
         xed_tables_init();
   }

   //Open Output file/pipe
   openOutput();

   //megabytes to bytes, go with power of ten for disk
   //will detach if larger than 
   fileLimit = FileLimit.Value() * 1000000;

   //Register Fini to be called when the application exits
   PIN_AddFiniFunction(Fini, NULL);
   //Registers a function that will be called if the pin exits early
   PIN_AddDetachFunction(DetachCallback, NULL);

   //ThreadStart is called when an application thread is created
   PIN_AddThreadStartFunction(ThreadStart, NULL);
   //used to delete local storage
   PIN_AddThreadFiniFunction(ThreadFini, NULL);
   
   //instrumentation modes
   if( ASCII.Value() )
   {
      if( AnnotateDisassembly.Value() ) 
         TRACE_AddInstrumentFunction(DisControlTrace, NULL);   
      else
         TRACE_AddInstrumentFunction(ASCIIControlTrace, NULL);    
   }
   else
      TRACE_AddInstrumentFunction(BinaryControlTrace, NULL);   

   //exception callback 
   PIN_AddContextChangeFunction(Sig, NULL);
   
   //Enable routine filter
   filter.Activate();
   // Start the program, never returns
   gettimeofday(&t1, NULL);
   PIN_StartProgram();
    
   return 0;
}
